#!/usr/bin/env python3
"""
统一配色方案

定义渐变配色方案和不同实验类型的配色映射
"""

from typing import Dict, List, Tuple, Optional
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors


class ColorPalette:
    """配色方案类"""
    
    # 主色调渐变
    BLUE_GRADIENT = ['#1f77b4', '#aec7e8']  # 深蓝 → 浅蓝
    ORANGE_GRADIENT = ['#ff7f0e', '#ffbb78']  # 深橙 → 浅橙
    GREEN_GRADIENT = ['#2ca02c', '#98df8a']  # 深绿 → 浅绿
    PURPLE_GRADIENT = ['#9467bd', '#c5b0d5']  # 深紫 → 浅紫
    RED_GRADIENT = ['#d62728', '#ff9896']  # 深红 → 浅红
    
    # 背景色
    BACKGROUND_WHITE = '#ffffff'
    BACKGROUND_LIGHT_GRAY = '#f5f5f5'
    
    # 专业热力图配色
    VIRIDIS = 'viridis'
    PLASMA = 'plasma'
    INFERNO = 'inferno'
    COOLWARM = 'coolwarm'
    BLUE_WHITE_RED = 'RdBu_r'
    
    def __init__(self):
        """初始化配色方案"""
        self.palettes = {
            'blue': self.BLUE_GRADIENT,
            'orange': self.ORANGE_GRADIENT,
            'green': self.GREEN_GRADIENT,
            'purple': self.PURPLE_GRADIENT,
            'red': self.RED_GRADIENT,
        }
    
    def get_gradient_colors(self, 
                           base_color: str,
                           n_colors: int = 10) -> List[str]:
        """
        生成渐变颜色列表
        
        Args:
            base_color: 基础颜色名称 ('blue', 'orange', 'green', 'purple', 'red')
            n_colors: 颜色数量
            
        Returns:
            渐变颜色列表
        """
        if base_color not in self.palettes:
            base_color = 'blue'
        
        gradient = self.palettes[base_color]
        cmap = mcolors.LinearSegmentedColormap.from_list(
            'custom', gradient, N=n_colors
        )
        colors = [mcolors.rgb2hex(cmap(i)) for i in np.linspace(0, 1, n_colors)]
        
        return colors
    
    def get_comparison_colors(self, n_methods: int) -> List[str]:
        """
        获取对比实验的颜色（不同色系）
        
        Args:
            n_methods: 方法数量
            
        Returns:
            颜色列表
        """
        color_names = ['blue', 'orange', 'green', 'purple', 'red']
        colors = []
        
        for i in range(n_methods):
            color_name = color_names[i % len(color_names)]
            gradient_colors = self.get_gradient_colors(color_name, 5)
            colors.append(gradient_colors[2])  # 取中间色
        
        return colors
    
    def get_ablation_colors(self, n_components: int) -> List[str]:
        """
        获取消融实验的颜色（同一色系不同深浅）
        
        Args:
            n_components: 组件数量
            
        Returns:
            颜色列表
        """
        gradient_colors = self.get_gradient_colors('blue', n_components)
        return gradient_colors
    
    def get_heatmap_colormap(self, style: str = 'viridis'):
        """
        获取热力图配色方案
        
        Args:
            style: 配色风格 ('viridis', 'plasma', 'coolwarm', 'RdBu_r')
            
        Returns:
            matplotlib colormap
        """
        colormap_map = {
            'viridis': plt.cm.viridis,
            'plasma': plt.cm.plasma,
            'inferno': plt.cm.inferno,
            'coolwarm': plt.cm.coolwarm,
            'RdBu_r': plt.cm.RdBu_r,
        }
        
        return colormap_map.get(style, plt.cm.viridis)


def get_color_palette() -> ColorPalette:
    """获取默认配色方案"""
    return ColorPalette()
